from typing import List, Dict
import torch
import os
import shutil
from torch.utils.tensorboard import SummaryWriter

from utils import Timer
from config import workspace, weight_root, metrics_root, tensorboard_root, key_path, trains, valids, config_backup_path, metric_csv_sep


class Logger(object):
    """
    用于训练过程中的日志记录
    """
    def __init__(
            self,
            model: torch.nn.Module,
            metrics_seg: List[str],
            metrics_cls: List[str],
            metrics_val: List[str],
            T: Timer = None
    ):
        self.model = model
        self.T = T
        self.ET = None
        self.name_set: Dict[str, List[str]] = {
            'seg': metrics_seg,      # segmentation
            'cls': metrics_cls,      # classify
            'etr': metrics_val,     # evaluate_train
            'evd': metrics_val,     # evaluate_valid
        }

    def start_epoch(self, epoch: int) -> None:
        if self.T: self.T.track(f' -> epoch {epoch} training start')
        if self.T: self.ET = self.T.tab()

    def log_epoch(self, prefix: str, items: Dict[str, float], epoch: int) -> None:
        for name, value in items.items():
            self.tensorboard.add_scalar(tag=f'{prefix}-{name}', scalar_value=value)
        if self.ET:
            line = [f"{name}:%.4f" % value for name, value in items.items()]
            self.ET.track(f' -> {prefix}: {line}')
        self.metrics[prefix].write(f'{str(epoch)},{metric_csv_sep.join(["%.4f" % items[name] for name in self.name_set[prefix]])}\n')
        self.metrics[prefix].flush()

    def save_epoch(self, epoch: int) -> None:
        if self.T: self.T.track(f' -> epoch {epoch} saving weights')
        torch.save(self.model, weight_root / f'auto-save-epoch-{epoch}.pth')

    def __enter__(self):
        if self.T: self.T.track(f' -> start training')
        # 创建相关目录
        os.makedirs(weight_root, exist_ok=True)
        os.makedirs(metrics_root, exist_ok=True)
        # 保存参数文件到本地
        shutil.copy(workspace / 'main' / 'config.py', config_backup_path)
        # 保存分组信息到本地
        with open(key_path, 'w') as f:
            f.writelines(f'trains: {trains}\n')
            f.writelines(f'valids: {valids}')
        # tensorboard 日志
        self.tensorboard = SummaryWriter(log_dir=tensorboard_root)
        # 自己的日志
        self.metrics = {prefix: open(metrics_root / f'{prefix}.txt', 'w+') for prefix in ['seg', 'cls', 'etr', 'evd']}
        # 日志写头行
        for prefix, metric in self.metrics.items():
            names = ['epoch'] + [f'{prefix}-{name}' for name in self.name_set[prefix]]
            metric.write(f'{metric_csv_sep.join(names)}\n')

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.T: self.T.track(f' -> ending training')
        self.tensorboard.close()
        for metric in self.metrics.values():
            metric.close()
        if exc_type or exc_val or exc_tb:
            if self.T: self.T.track(f' -> exc_type: {exc_type}')
            if self.T: self.T.track(f' -> exc_value: {exc_val}')
            if self.T: self.T.track(f' -> exc_traceback: {exc_tb}')
        return False
